import torch
import torch.nn as nn
import torch.nn.functional as F
from modules.yolov3.yololayer import YoloLayer

# <function parse_model_config>
def parse_model_config(path):
    """Parses the yolo-v3 layer configuration file and returns module definitions"""
    file = open(path, 'r')
    lines = file.read().split('\n')
    lines = [x for x in lines if x and not x.startswith('#')]
    lines = [x.rstrip().lstrip() for x in lines] # get rid of fringe whitespaces
    module_defs = []
    for line in lines:
        if line.startswith('['): # This marks the start of a new block
            module_defs.append({})
            module_defs[-1]['type'] = line[1:-1].rstrip()
            if module_defs[-1]['type'] == 'convolutional':
                module_defs[-1]['batch_normalize'] = 0
        else:
            key, value = line.split("=")
            value = value.strip()
            module_defs[-1][key.rstrip()] = value.strip()
        # end-if
    # end-for
    return module_defs
# <function parse_model_config>

# <function parse_data_config>
def parse_data_config(path):
    """Parses the data configuration file"""
    options = dict()
    options['gpus'] = '0,1,2,3'
    options['num_workers'] = '10'
    with open(path, 'r') as fp:
        lines = fp.readlines()
    # end-with
    for line in lines:
        line = line.strip()
        if line == '' or line.startswith('#'):
            continue
        # end-if
        key, value = line.split('=')
        options[key.strip()] = value.strip()
    # end-for
    return options
# <function parse_data_config>

# <function create_modules>
def create_modules(module_defs):
    """
    Constructs module list of layer blocks from module configuration in module_defs
    """
    hyperparams = module_defs.pop(0)
    output_filters = [int(hyperparams["channels"])]
    module_list = nn.ModuleList()
    for module_i, module_def in enumerate(module_defs):
        modules = nn.Sequential()
        pass
        if module_def["type"] == "convolutional":
            bn = int(module_def["batch_normalize"])
            filters = int(module_def["filters"])
            kernel_size = int(module_def["size"])
            pad = (kernel_size - 1) // 2
            modules.add_module(
                f"conv_{module_i}",
                nn.Conv2d(
                    in_channels=output_filters[-1],
                    out_channels=filters,
                    kernel_size=kernel_size,
                    stride=int(module_def["stride"]),
                    padding=pad,
                    bias=not bn,
                ),
            )
            if bn:
                modules.add_module(f"batch_norm_{module_i}", nn.BatchNorm2d(filters, momentum=0.9, eps=1e-5))
            # end-if
            if module_def["activation"] == "leaky":
                modules.add_module(f"leaky_{module_i}", nn.LeakyReLU(0.1))
            # end-if
        elif module_def["type"] == "maxpool":
            kernel_size = int(module_def["size"])
            stride = int(module_def["stride"])
            if kernel_size == 2 and stride == 1:
                modules.add_module(f"_debug_padding_{module_i}", nn.ZeroPad2d((0, 1, 0, 1)))
            maxpool = nn.MaxPool2d(kernel_size=kernel_size, stride=stride, padding=int((kernel_size - 1) // 2))
            modules.add_module(f"maxpool_{module_i}", maxpool)
            pass
        elif module_def["type"] == "upsample":
            upsample = Upsample(scale_factor=int(module_def["stride"]), mode="nearest")
            modules.add_module(f"upsample_{module_i}", upsample)
            pass
        elif module_def["type"] == "route":
            layers = [int(x) for x in module_def["layers"].split(",")]
            filters = sum([output_filters[1:][i] for i in layers])
            modules.add_module(f"route_{module_i}", EmptyLayer())
            pass
        elif module_def["type"] == "shortcut":
            filters = output_filters[1:][int(module_def["from"])]
            modules.add_module(f"shortcut_{module_i}", EmptyLayer())
            pass
        elif module_def["type"] == "yolo":
            modules.add_module(f"yolo_inp_{module_i}", EmptyLayer())
        # end-if
        # Register module list and number of output filters
        module_list.append(modules)
        output_filters.append(filters)
    # end-for
    return hyperparams, module_list
# <function create_modules>

# <class Upsample(nn.Module)>
class Upsample(nn.Module):
    """ nn.Upsample is deprecated """
    # <method __init__>
    def __init__(self, scale_factor, mode="nearest"):
        super(Upsample, self).__init__()
        self.scale_factor = scale_factor
        self.mode = mode
    # <method __init__>
    # <method forward>
    def forward(self, x):
        x = F.interpolate(x, scale_factor=self.scale_factor, mode=self.mode)
        return x
    # <method forward>
# <class Upsample(nn.Module)>

# <class EmptyLayer(nn.Module)>
class EmptyLayer(nn.Module):
    """Placeholder for 'route' and 'shortcut' layers"""
    # <method __init__>
    def __init__(self):
        super(EmptyLayer, self).__init__()
    # <method __init__>
    # <method forward>
    def forward(self, x):
        return x
    # <method forward>
# <class EmptyLayer(nn.Module)>

# <class DarkNetFeature(nn.Module)>
class DarkNetFeature(nn.Module):
    """YOLOv3 object detection model"""
    # <method __init__>
    def __init__(self, config_path):
        super(DarkNetFeature, self).__init__()
        self.module_defs = parse_model_config(config_path)
        self.hyperparams, self.module_list = create_modules(self.module_defs)
    # <method __init__>
    # <method forward>
    def forward(self, x):
        img_dim = x.shape[2]
        layer_outputs, yolo_outputs = [], []
        for i, (module_def, module) in enumerate(zip(self.module_defs, self.module_list)):
            if module_def["type"] in ["convolutional", "upsample", "maxpool"]:
                x = module(x)
            elif module_def["type"] == "route":
                x = torch.cat([layer_outputs[int(layer_i)] for layer_i in module_def["layers"].split(",")], 1)
            elif module_def["type"] == "shortcut":
                layer_i = int(module_def["from"])
                x = layer_outputs[-1] + layer_outputs[layer_i]
            elif module_def["type"] == "yolo":          
                yolo_outputs.append(x)
            # end-if
            layer_outputs.append(x)
        # end-for
        return yolo_outputs
    # <method forward>
# <class DarkNetFeature(nn.Module)>

# <class DarkNet(nn.Module)>
class DarkNet(nn.Module):
    # <method __init__>
    def __init__(
        self,         
        anchors,
        num_classes,
        config_path="./cfgs/yolov3-face.cfg"
        ):
        super(DarkNet, self).__init__()
        self._feature = DarkNetFeature(config_path=config_path)
        self._yolo_layers = []
        for idx in range(len(anchors)):
            self._yolo_layers.append( YoloLayer(anchors = anchors[idx], num_classes = num_classes) )
        pass
    # <method __init__>
    # <method forward>
    def forward(self, x):
        nI_h = x.size(2)
        nI_w = x.size(3)
        features = self._feature(x)
        encoded_fms = [ ]
        decoded_fms = [ ]
        for idx in range(len(features)):
            encoded_fm, decode_fm = self._yolo_layers[idx](features[idx], nI_h, nI_w)
            encoded_fms.append(encoded_fm)
            decoded_fms.append(decode_fm)
        # end-for
        return torch.cat(encoded_fms, 1), torch.cat(decoded_fms, 1)
    # <method forward>
# <class DarkNet(nn.Module)>

if __name__ == "__main__":
    anchors = (
        ((116, 90), (156, 198), (373, 326)),
        ((30,  61), (62,   45), (59,  119)), 
        ((10,  13), (16,   30), (33,  23)), 
        )
    inp = torch.rand(1, 3, 416, 416)        
    net = DarkNet(anchors, 2, "./cfgs/yolov3-face.cfg")
    outp = net(inp)
    import pdb; pdb.set_trace()